package com.glsc.ngateway.common.base.utils;

import org.apache.commons.collections4.CollectionUtils;

import javax.validation.constraints.NotNull;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.CountDownLatch;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Stream;

public class TreeNodeUtils {

    /**
     * BFS遍历
     *
     * @param root
     * @param childrenFunc
     * @param <T>
     * @return
     */
    public static <T> List<T> flatTree(@NotNull T root, @NotNull Function<T, List<T>> childrenFunc) {
        Queue<T> que = new LinkedList<>();
        List<T> res = new ArrayList<>();
        que.add(root);
        while (!que.isEmpty()) {
            T node = que.poll();
            res.add(node);
            List<T> children = childrenFunc.apply(node);
            if (CollectionUtils.isNotEmpty(children)) {
                que.addAll(children);
            }
        }
        return res;
    }


    public static <T> void BFSWalk(@NotNull T root, @NotNull Function<T, List<T>> childrenFunc, Consumer<T> func) throws InterruptedException {
        BFSWalk(root, childrenFunc, func, false);
    }

    /**
     *
     * @param root
     * @param childrenFunc
     * @param func
     * @param parallelStream 是否开启并行流，在每一层中，并行执行func
     * @param <T>
     */
    public static <T> void BFSWalk(@NotNull T root, @NotNull Function<T, List<T>> childrenFunc, Consumer<T> func, boolean parallelStream) throws InterruptedException {
        Queue<T> que = new LinkedList<>();
        que.add(root);
        while (!que.isEmpty()) {
            CountDownLatch count = new CountDownLatch(que.size());
            List<T> nodeList = new ArrayList<>();
            while (!que.isEmpty()) {
                nodeList.add(que.poll());
            }
            Stream<T> stream = parallelStream ? nodeList.parallelStream(): nodeList.stream();
            stream.forEach(node -> {
                //每层多线程
                func.accept(node);
                count.countDown();
                List<T> children = childrenFunc.apply(node);
                if (CollectionUtils.isNotEmpty(children)) {
                    que.addAll(children);
                }
            });
            count.await();
        }
    }
}
